import pandas as pd
import mplfinance as mpf


def plot_k_line(datas, png_url):
    try:
        # 将数据集转换为 DataFrame
        df = pd.DataFrame(datas)
        # 将日期列转换为 pandas 的日期时间类型
        # df['Date'] = pd.to_datetime(df['Date'])
        # 将日期列设置为索引
        df.set_index('timekey', inplace=True)

        # 绘制 K 线图
        mpf.plot(df, type='candle', style='yahoo', savefig=png_url)
        print(f"K 线图已成功保存到 {png_url}")
    except Exception as e:
        print(f"绘制 K 线图时出现错误: {e}")

if __name__=="__main__":
    # 示例数据集
    datas = [
        {'Date': '2023-01-01', 'Open': 100, 'High': 105, 'Low': 95, 'Close': 102},
        {'Date': '2023-01-02', 'Open': 102, 'High': 108, 'Low': 101, 'Close': 106},
        {'Date': '2023-01-03', 'Open': 106, 'High': 107, 'Low': 103, 'Close': 104}
    ]

    # 指定保存图片的路径
    png_url = 'k_line_chart.png'

    # 调用函数绘制并保存 K 线图
    plot_k_line(datas, png_url)
